# -*- coding: utf-8 -*-
# @File     : voc_xml_to_txt.py
# @Author   : bingjia
# @Time     : 2020/7/31 21:24
# @Desc     : 将voc格式的xml文件转为txt格式

import os
import shutil
import random
import xml.etree.ElementTree as ET

xml_path = "../data/xmls"
classes = ['flora', 'shit', 'break', 'dust', 'strip']  # 修改为自己的类别


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh

    return (x, y, w, h)


def convert_annotation(xml_id):
    """
    将精灵标注的xml文件转为txt文件
    :param xml_id:
    :return:
    """
    in_file = open(os.path.join(os.path.dirname(os.getcwd()), 'data', 'xmls', '%s.xml' % (xml_id)), encoding='utf-8')
    out_file = open(os.path.join(os.path.dirname(os.getcwd()), 'data', 'txts', '%s.txt' % (xml_id)), 'w')

    tree = ET.parse(in_file)
    root = tree.getroot()

    try:
        size = root.find('size')
        w, h = int(size.find('width').text), int(size.find('height').text)

        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes or int(difficult) == 1:
                continue
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            bb = convert((w, h), b)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    except Exception as e:
        print("该文件为负样本...")
    else:
        pass


# for xml in os.listdir(os.path.join(os.path.dirname(os.getcwd()), 'data', 'xmls')):
#     print(xml)
#     convert_annotation(xml.split(".")[0])


def data_dispatch(image_path, txt_path, ratio=0.8):
    """
    将图片数据和标签数据按照比例分开(train\val)
    :param image_path:
    :param txt_path:
    :param ratio: 比例
    :return:
    """
    data_list = [str(data.split(".")[0]) for data in os.listdir(os.path.join(txt_path))]
    val_data_list = random.sample(data_list, int((1 - ratio) * len(data_list)))
    train_data_list = list(set(data_list).difference(set(val_data_list)))

    for i in data_list:
        print(i)
        txt_name = os.path.join(txt_path, i + ".txt")
        image_name = os.path.join(image_path, i + ".jpg")

        if i in val_data_list:
            shutil.copy(txt_name,
                        os.path.join(os.path.dirname(os.getcwd()), 'transformer', 'labels', 'val', i + ".txt"))
            shutil.copy(image_name,
                        os.path.join(os.path.dirname(os.getcwd()), 'transformer', 'images', 'val', i + ".jpg"))
        elif i in train_data_list:
            shutil.copy(txt_name,
                        os.path.join(os.path.dirname(os.getcwd()), 'transformer', 'labels', 'train', i + ".txt"))
            shutil.copy(image_name,
                        os.path.join(os.path.dirname(os.getcwd()), 'transformer', 'images', 'train', i + ".jpg"))


# data_dispatch(
#     txt_path=r'D:\AppStore\PycharmStore\operational_master\operation_back\operation\trains\network\yolov5\data\txts',
#     image_path=r'D:\AppStore\PycharmStore\operational_master\operation_back\operation\trains\network\yolov5\data\images',
#     ratio=0.8)
